#include "skynet.h"
#include "skynet_harbor.h"
#include "skynet_socket.h"
#include "skynet_handle.h"

/*
        harbor listen the PTYPE_HARBOR (in text)
        N name : update the global name
        S fd id: connect to new harbor , we should send self_id to fd first ,
   and then recv a id (check it), and at last send queue. A fd id: accept new
   harbor , we should send self_id to fd , and then send queue.

        If the fd is disconnected, send message to slave in PTYPE_TEXT.  D id
        If we don't known a globalname, send message to slave in PTYPE_TEXT. Q
   name
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <assert.h>
#include <stdint.h>
#include <unistd.h>

#define HASH_SIZE 4096
#define DEFAULT_QUEUE_SIZE 1024

// 12 is sizeof(struct remote_message_header)
#define HEADER_COOKIE_LENGTH 12

/*
        message type (8bits) is in destination high 8bits
        harbor id (8bits) is also in that place , but remote message doesn't
   need harbor id.
 */
struct remote_message_header {
  uint32_t source;
  uint32_t destination;
  uint32_t session;
};

struct harbor_msg {
  struct remote_message_header header;
  void *buffer;
  size_t size;
};

struct harbor_msg_queue {
  int size;
  int head;
  int tail;
  struct harbor_msg *data;
};

struct keyvalue {
  struct keyvalue *next;
  char key[GLOBALNAME_LENGTH];
  uint32_t hash;
  uint32_t value;
  struct harbor_msg_queue *queue;
};

struct hashmap {
  struct keyvalue *node[HASH_SIZE];
};

#define STATUS_WAIT 0
#define STATUS_HANDSHAKE 1
#define STATUS_HEADER 2
#define STATUS_CONTENT 3
#define STATUS_DOWN 4

struct slave {
  int fd;
  struct harbor_msg_queue *queue;
  int status;
  int length;
  int read;
  uint8_t size[4];
  char *recv_buffer;
};

struct harbor {
  struct skynet_context *ctx;
  int id;
  uint32_t slave;
  struct hashmap *map;
  struct slave s[REMOTE_MAX];
};

// hash table

static void push_queue_msg(struct harbor_msg_queue *queue,
                           struct harbor_msg *m) {
  // If there is only 1 free slot which is reserved to distinguish full/empty
  // of circular buffer, expand it.
  if (((queue->tail + 1) % queue->size) == queue->head) {
    struct harbor_msg *new_buffer =
        skynet_malloc(queue->size * 2 * sizeof(struct harbor_msg));
    int i;
    for (i = 0; i < queue->size - 1; i++) {
      new_buffer[i] = queue->data[(i + queue->head) % queue->size];
    }
    skynet_free(queue->data);
    queue->data = new_buffer;
    queue->head = 0;
    queue->tail = queue->size - 1;
    queue->size *= 2;
  }
  struct harbor_msg *slot = &queue->data[queue->tail];
  *slot = *m;
  queue->tail = (queue->tail + 1) % queue->size;
}

static void push_queue(struct harbor_msg_queue *queue, void *buffer, size_t sz,
                       struct remote_message_header *header) {
  struct harbor_msg m;
  m.header = *header;
  m.buffer = buffer;
  m.size = sz;
  push_queue_msg(queue, &m);
}

static struct harbor_msg *pop_queue(struct harbor_msg_queue *queue) {
  if (queue->head == queue->tail) {
    return NULL;
  }
  struct harbor_msg *slot = &queue->data[queue->head];
  queue->head = (queue->head + 1) % queue->size;
  return slot;
}

static struct harbor_msg_queue *new_queue() {
  struct harbor_msg_queue *queue = skynet_malloc(sizeof(*queue));
  queue->size = DEFAULT_QUEUE_SIZE;
  queue->head = 0;
  queue->tail = 0;
  queue->data = skynet_malloc(DEFAULT_QUEUE_SIZE * sizeof(struct harbor_msg));

  return queue;
}

static void release_queue(struct harbor_msg_queue *queue) {
  if (queue == NULL)
    return;
  struct harbor_msg *m;
  while ((m = pop_queue(queue)) != NULL) {
    skynet_free(m->buffer);
  }
  skynet_free(queue->data);
  skynet_free(queue);
}

static struct keyvalue *hash_search(struct hashmap *hash,
                                    const char name[GLOBALNAME_LENGTH]) {
  uint32_t *ptr = (uint32_t *)name;
  uint32_t h = ptr[0] ^ ptr[1] ^ ptr[2] ^ ptr[3];
  struct keyvalue *node = hash->node[h % HASH_SIZE];
  while (node) {
    if (node->hash == h && strncmp(node->key, name, GLOBALNAME_LENGTH) == 0) {
      return node;
    }
    node = node->next;
  }
  return NULL;
}

/*

// Don't support erase name yet

static struct void
hash_erase(struct hashmap * hash, char name[GLOBALNAME_LENGTH) {
        uint32_t *ptr = name;
        uint32_t h = ptr[0] ^ ptr[1] ^ ptr[2] ^ ptr[3];
        struct keyvalue ** ptr = &hash->node[h % HASH_SIZE];
        while (*ptr) {
                struct keyvalue * node = *ptr;
                if (node->hash == h && strncmp(node->key, name,
GLOBALNAME_LENGTH) == 0) { _release_queue(node->queue); *ptr->next = node->next;
                        skynet_free(node);
                        return;
                }
                *ptr = &(node->next);
        }
}
*/

static struct keyvalue *hash_insert(struct hashmap *hash,
                                    const char name[GLOBALNAME_LENGTH]) {
  uint32_t *ptr = (uint32_t *)name;
  uint32_t h = ptr[0] ^ ptr[1] ^ ptr[2] ^ ptr[3];
  struct keyvalue **pkv = &hash->node[h % HASH_SIZE];
  struct keyvalue *node = skynet_malloc(sizeof(*node));
  memcpy(node->key, name, GLOBALNAME_LENGTH);
  node->next = *pkv;
  node->queue = NULL;
  node->hash = h;
  node->value = 0;
  *pkv = node;

  return node;
}

static struct hashmap *hash_new() {
  struct hashmap *h = skynet_malloc(sizeof(struct hashmap));
  memset(h, 0, sizeof(*h));
  return h;
}

static void hash_delete(struct hashmap *hash) {
  int i;
  for (i = 0; i < HASH_SIZE; i++) {
    struct keyvalue *node = hash->node[i];
    while (node) {
      struct keyvalue *next = node->next;
      release_queue(node->queue);
      skynet_free(node);
      node = next;
    }
  }
  skynet_free(hash);
}

///////////////

static void close_harbor(struct harbor *h, int id) {
  struct slave *s = &h->s[id];
  s->status = STATUS_DOWN;
  if (s->fd) {
    skynet_socket_close(h->ctx, s->fd);
    s->fd = 0;
  }
  if (s->queue) {
    release_queue(s->queue);
    s->queue = NULL;
  }
}

static void report_harbor_down(struct harbor *h, int id) {
  char down[64];
  int n = sprintf(down, "D %d", id);

  skynet_send(h->ctx, 0, h->slave, PTYPE_TEXT, 0, down, n);
}

struct harbor *harbor_create(void) {
  struct harbor *h = skynet_malloc(sizeof(*h));
  memset(h, 0, sizeof(*h));
  h->map = hash_new();
  return h;
}

static void close_all_remotes(struct harbor *h) {
  int i;
  for (i = 1; i < REMOTE_MAX; i++) {
    close_harbor(h, i);
    // don't call report_harbor_down.
    // never call skynet_send during module exit, because of dead lock
  }
}

void harbor_release(struct harbor *h) {
  close_all_remotes(h);
  hash_delete(h->map);
  skynet_free(h);
}

static inline void to_bigendian(uint8_t *buffer, uint32_t n) {
  buffer[0] = (n >> 24) & 0xff;
  buffer[1] = (n >> 16) & 0xff;
  buffer[2] = (n >> 8) & 0xff;
  buffer[3] = n & 0xff;
}

static inline void header_to_message(const struct remote_message_header *header,
                                     uint8_t *message) {
  to_bigendian(message, header->source);
  to_bigendian(message + 4, header->destination);
  to_bigendian(message + 8, header->session);
}

static inline uint32_t from_bigendian(uint32_t n) {
  union {
    uint32_t big;
    uint8_t bytes[4];
  } u;
  u.big = n;
  return u.bytes[0] << 24 | u.bytes[1] << 16 | u.bytes[2] << 8 | u.bytes[3];
}

static inline void message_to_header(const uint32_t *message,
                                     struct remote_message_header *header) {
  header->source = from_bigendian(message[0]);
  header->destination = from_bigendian(message[1]);
  header->session = from_bigendian(message[2]);
}

// socket package

static void forward_local_messsage(struct harbor *h, void *msg, int sz) {
  const char *cookie = msg;
  cookie += sz - HEADER_COOKIE_LENGTH;
  struct remote_message_header header;
  message_to_header((const uint32_t *)cookie, &header);

  uint32_t destination = header.destination;
  int type = destination >> HANDLE_REMOTE_SHIFT;
  destination =
      (destination & HANDLE_MASK) | ((uint32_t)h->id << HANDLE_REMOTE_SHIFT);

  if (skynet_send(h->ctx, header.source, destination, type | PTYPE_TAG_DONTCOPY,
                  (int)header.session, (void *)msg,
                  sz - HEADER_COOKIE_LENGTH) < 0) {
    if (type != PTYPE_ERROR) {
      // don't need report error when type is error
      skynet_send(h->ctx, destination, header.source, PTYPE_ERROR,
                  (int)header.session, NULL, 0);
    }
    skynet_error(h->ctx, "Unknown destination :%x from :%x type(%d)",
                 destination, header.source, type);
  }
}

static void send_remote(struct skynet_context *ctx, int fd, const char *buffer,
                        size_t sz, struct remote_message_header *cookie) {
  size_t sz_header = sz + sizeof(*cookie);
  if (sz_header > UINT32_MAX) {
    skynet_error(ctx, "remote message from :%08x to :%08x is too large.",
                 cookie->source, cookie->destination);
    return;
  }
  uint8_t sendbuf[sz_header + 4];
  to_bigendian(sendbuf, (uint32_t)sz_header);
  memcpy(sendbuf + 4, buffer, sz);
  header_to_message(cookie, sendbuf + 4 + sz);

  struct socket_sendbuffer tmp;
  tmp.id = fd;
  tmp.type = SOCKET_BUFFER_RAWPOINTER;
  tmp.buffer = sendbuf;
  tmp.sz = sz_header + 4;

  // ignore send error, because if the connection is broken, the mainloop will
  // recv a message.
  skynet_socket_sendbuffer(ctx, &tmp);
}

static void dispatch_name_queue(struct harbor *h, struct keyvalue *node) {
  struct harbor_msg_queue *queue = node->queue;
  uint32_t handle = node->value;
  int harbor_id = handle >> HANDLE_REMOTE_SHIFT;
  struct skynet_context *context = h->ctx;
  struct slave *s = &h->s[harbor_id];
  int fd = s->fd;
  if (fd == 0) {
    if (s->status == STATUS_DOWN) {
      char tmp[GLOBALNAME_LENGTH + 1];
      memcpy(tmp, node->key, GLOBALNAME_LENGTH);
      tmp[GLOBALNAME_LENGTH] = '\0';
      skynet_error(context, "Drop message to %s (in harbor %d)", tmp,
                   harbor_id);
    } else {
      if (s->queue == NULL) {
        s->queue = node->queue;
        node->queue = NULL;
      } else {
        struct harbor_msg *m;
        while ((m = pop_queue(queue)) != NULL) {
          push_queue_msg(s->queue, m);
        }
      }
      if (harbor_id == (h->slave >> HANDLE_REMOTE_SHIFT)) {
        // the harbor_id is local
        struct harbor_msg *m;
        while ((m = pop_queue(s->queue)) != NULL) {
          int type = m->header.destination >> HANDLE_REMOTE_SHIFT;
          skynet_send(context, m->header.source, handle,
                      type | PTYPE_TAG_DONTCOPY, m->header.session, m->buffer,
                      m->size);
        }
        release_queue(s->queue);
        s->queue = NULL;
      }
    }
    return;
  }
  struct harbor_msg *m;
  while ((m = pop_queue(queue)) != NULL) {
    m->header.destination |= (handle & HANDLE_MASK);
    send_remote(context, fd, m->buffer, m->size, &m->header);
    skynet_free(m->buffer);
  }
}

static void dispatch_queue(struct harbor *h, int id) {
  struct slave *s = &h->s[id];
  int fd = s->fd;
  assert(fd != 0);

  struct harbor_msg_queue *queue = s->queue;
  if (queue == NULL)
    return;

  struct harbor_msg *m;
  while ((m = pop_queue(queue)) != NULL) {
    send_remote(h->ctx, fd, m->buffer, m->size, &m->header);
    skynet_free(m->buffer);
  }
  release_queue(queue);
  s->queue = NULL;
}

static void push_socket_data(struct harbor *h,
                             const struct skynet_socket_message *message) {
  assert(message->type == SKYNET_SOCKET_TYPE_DATA);
  int fd = message->id;
  int i;
  int id = 0;
  struct slave *s = NULL;
  for (i = 1; i < REMOTE_MAX; i++) {
    if (h->s[i].fd == fd) {
      s = &h->s[i];
      id = i;
      break;
    }
  }
  if (s == NULL) {
    skynet_error(h->ctx, "Invalid socket fd (%d) data", fd);
    return;
  }
  uint8_t *buffer = (uint8_t *)message->buffer;
  int size = message->ud;

  for (;;) {
    switch (s->status) {
    case STATUS_HANDSHAKE: {
      // check id
      uint8_t remote_id = buffer[0];
      if (remote_id != id) {
        skynet_error(h->ctx,
                     "Invalid shakehand id (%d) from fd = %d , harbor = %d", id,
                     fd, remote_id);
        close_harbor(h, id);
        return;
      }
      ++buffer;
      --size;
      s->status = STATUS_HEADER;

      dispatch_queue(h, id);

      if (size == 0) {
        break;
      }
      // go though
    }
    case STATUS_HEADER: {
      // big endian 4 bytes length, the first one must be 0.
      int need = 4 - s->read;
      if (size < need) {
        memcpy(s->size + s->read, buffer, size);
        s->read += size;
        return;
      } else {
        memcpy(s->size + s->read, buffer, need);
        buffer += need;
        size -= need;

        if (s->size[0] != 0) {
          skynet_error(h->ctx, "Message is too long from harbor %d", id);
          close_harbor(h, id);
          return;
        }
        s->length = s->size[1] << 16 | s->size[2] << 8 | s->size[3];
        s->read = 0;
        s->recv_buffer = skynet_malloc(s->length);
        s->status = STATUS_CONTENT;
        if (size == 0) {
          return;
        }
      }
    }
    // go though
    case STATUS_CONTENT: {
      int need = s->length - s->read;
      if (size < need) {
        memcpy(s->recv_buffer + s->read, buffer, size);
        s->read += size;
        return;
      }
      memcpy(s->recv_buffer + s->read, buffer, need);
      forward_local_messsage(h, s->recv_buffer, s->length);
      s->length = 0;
      s->read = 0;
      s->recv_buffer = NULL;
      size -= need;
      buffer += need;
      s->status = STATUS_HEADER;
      if (size == 0)
        return;
      break;
    }
    default:
      return;
    }
  }
}

static void update_name(struct harbor *h, const char name[GLOBALNAME_LENGTH],
                        uint32_t handle) {
  struct keyvalue *node = hash_search(h->map, name);
  if (node == NULL) {
    node = hash_insert(h->map, name);
  }
  node->value = handle;
  if (node->queue) {
    dispatch_name_queue(h, node);
    release_queue(node->queue);
    node->queue = NULL;
  }
}

static int remote_send_handle(struct harbor *h, uint32_t source,
                              uint32_t destination, int type, int session,
                              const char *msg, size_t sz) {
  int harbor_id = destination >> HANDLE_REMOTE_SHIFT;
  struct skynet_context *context = h->ctx;
  if (harbor_id == h->id) {
    // local message
    skynet_send(context, source, destination, type | PTYPE_TAG_DONTCOPY,
                session, (void *)msg, sz);
    return 1;
  }

  struct slave *s = &h->s[harbor_id];
  if (s->fd == 0 || s->status == STATUS_HANDSHAKE) {
    if (s->status == STATUS_DOWN) {
      // throw an error return to source
      // report the destination is dead
      skynet_send(context, destination, source, PTYPE_ERROR, session, NULL, 0);
      skynet_error(
          context,
          "Drop message to harbor %d from %x to %x (session = %d, msgsz = %d)",
          harbor_id, source, destination, session, (int)sz);
    } else {
      if (s->queue == NULL) {
        s->queue = new_queue();
      }
      struct remote_message_header header;
      header.source = source;
      header.destination =
          (type << HANDLE_REMOTE_SHIFT) | (destination & HANDLE_MASK);
      header.session = (uint32_t)session;
      push_queue(s->queue, (void *)msg, sz, &header);
      return 1;
    }
  } else {
    struct remote_message_header cookie;
    cookie.source = source;
    cookie.destination =
        (destination & HANDLE_MASK) | ((uint32_t)type << HANDLE_REMOTE_SHIFT);
    cookie.session = (uint32_t)session;
    send_remote(context, s->fd, msg, sz, &cookie);
  }

  return 0;
}

static int remote_send_name(struct harbor *h, uint32_t source,
                            const char name[GLOBALNAME_LENGTH], int type,
                            int session, const char *msg, size_t sz) {
  struct keyvalue *node = hash_search(h->map, name);
  if (node == NULL) {
    node = hash_insert(h->map, name);
  }
  if (node->value == 0) {
    if (node->queue == NULL) {
      node->queue = new_queue();
    }
    struct remote_message_header header;
    header.source = source;
    header.destination = type << HANDLE_REMOTE_SHIFT;
    header.session = (uint32_t)session;
    push_queue(node->queue, (void *)msg, sz, &header);
    char query[2 + GLOBALNAME_LENGTH + 1] = "Q ";
    query[2 + GLOBALNAME_LENGTH] = 0;
    memcpy(query + 2, name, GLOBALNAME_LENGTH);
    skynet_send(h->ctx, 0, h->slave, PTYPE_TEXT, 0, query, strlen(query));
    return 1;
  } else {
    return remote_send_handle(h, source, node->value, type, session, msg, sz);
  }
}

static void handshake(struct harbor *h, int id) {
  struct slave *s = &h->s[id];
  uint8_t handshake[1] = {(uint8_t)h->id};
  struct socket_sendbuffer tmp;
  tmp.id = s->fd;
  tmp.type = SOCKET_BUFFER_RAWPOINTER;
  tmp.buffer = handshake;
  tmp.sz = 1;
  skynet_socket_sendbuffer(h->ctx, &tmp);
}

static void harbor_command(struct harbor *h, const char *msg, size_t sz,
                           int session, uint32_t source) {
  const char *name = msg + 2;
  int s = (int)sz;
  s -= 2;
  switch (msg[0]) {
  case 'N': {
    if (s <= 0 || s >= GLOBALNAME_LENGTH) {
      skynet_error(h->ctx, "Invalid global name %s", name);
      return;
    }
    struct remote_name rn;
    memset(&rn, 0, sizeof(rn));
    memcpy(rn.name, name, s);
    rn.handle = source;
    update_name(h, rn.name, rn.handle);
    break;
  }
  case 'S':
  case 'A': {
    char buffer[s + 1];
    memcpy(buffer, name, s);
    buffer[s] = 0;
    int fd = 0, id = 0;
    sscanf(buffer, "%d %d", &fd, &id);
    if (fd == 0 || id <= 0 || id >= REMOTE_MAX) {
      skynet_error(h->ctx, "Invalid command %c %s", msg[0], buffer);
      return;
    }
    struct slave *slave = &h->s[id];
    if (slave->fd != 0) {
      skynet_error(h->ctx, "Harbor %d alreay exist", id);
      return;
    }
    slave->fd = fd;

    skynet_socket_start(h->ctx, fd);
    handshake(h, id);
    if (msg[0] == 'S') {
      slave->status = STATUS_HANDSHAKE;
    } else {
      slave->status = STATUS_HEADER;
      dispatch_queue(h, id);
    }
    break;
  }
  default:
    skynet_error(h->ctx, "Unknown command %s", msg);
    return;
  }
}

static int harbor_id(struct harbor *h, int fd) {
  int i;
  for (i = 1; i < REMOTE_MAX; i++) {
    struct slave *s = &h->s[i];
    if (s->fd == fd) {
      return i;
    }
  }
  return 0;
}

static int mainloop(struct skynet_context *context, void *ud, int type,
                    int session, uint32_t source, const void *msg, size_t sz) {
  struct harbor *h = ud;
  switch (type) {
  case PTYPE_SOCKET: {
    const struct skynet_socket_message *message = msg;
    switch (message->type) {
    case SKYNET_SOCKET_TYPE_DATA:
      push_socket_data(h, message);
      skynet_free(message->buffer);
      break;
    case SKYNET_SOCKET_TYPE_ERROR:
    case SKYNET_SOCKET_TYPE_CLOSE: {
      int id = harbor_id(h, message->id);
      if (id) {
        report_harbor_down(h, id);
      } else {
        skynet_error(context, "Unkown fd (%d) closed", message->id);
      }
      break;
    }
    case SKYNET_SOCKET_TYPE_CONNECT:
      // fd forward to this service
      break;
    case SKYNET_SOCKET_TYPE_WARNING: {
      int id = harbor_id(h, message->id);
      if (id) {
        skynet_error(context, "message havn't send to Harbor (%d) reach %d K",
                     id, message->ud);
      }
      break;
    }
    default:
      skynet_error(context, "recv invalid socket message type %d", type);
      break;
    }
    return 0;
  }
  case PTYPE_HARBOR: {
    harbor_command(h, msg, sz, session, source);
    return 0;
  }
  case PTYPE_SYSTEM: {
    // remote message out
    const struct remote_message *rmsg = msg;
    if (rmsg->destination.handle == 0) {
      if (remote_send_name(h, source, rmsg->destination.name, rmsg->type,
                           session, rmsg->message, rmsg->sz)) {
        return 0;
      }
    } else {
      if (remote_send_handle(h, source, rmsg->destination.handle, rmsg->type,
                             session, rmsg->message, rmsg->sz)) {
        return 0;
      }
    }
    skynet_free((void *)rmsg->message);
    return 0;
  }
  default:
    skynet_error(context, "recv invalid message from %x,  type = %d", source,
                 type);
    if (session != 0 && type != PTYPE_ERROR) {
      skynet_send(context, 0, source, PTYPE_ERROR, session, NULL, 0);
    }
    return 0;
  }
}

int harbor_init(struct harbor *h, struct skynet_context *ctx,
                const char *args) {
  h->ctx = ctx;
  int harbor_id = 0;
  uint32_t slave = 0;
  sscanf(args, "%d %u", &harbor_id, &slave);
  if (slave == 0) {
    return 1;
  }
  h->id = harbor_id;
  h->slave = slave;
  if (harbor_id == 0) {
    close_all_remotes(h);
  }
  skynet_callback(ctx, h, mainloop);
  skynet_harbor_start(ctx);

  return 0;
}
